import torch.nn as nn
import torch
from modules.models.backbones.mit import MixVisionTransformer
import torch.nn.functional as F

class SCTHead(nn.Module):
    def __init__(self,in_channels,feats_channels,num_classes,dropout_ratio=0.1, align_corners=False):
        super(SCTHead,self).__init__()
        self.bn1 = nn.SyncBatchNorm(in_channels)
        self.conv1 = nn.Conv2d(
            in_channels,
            feats_channels,
            kernel_size=3,
            padding=1)
        self.bn2 = nn.SyncBatchNorm(feats_channels)
        self.relu = nn.ReLU()
        if dropout_ratio > 0:
            self.dropout = nn.Dropout2d(dropout_ratio)
        else:
            self.dropout = None

        self.conv_seg = nn.Conv2d(feats_channels, num_classes, kernel_size=1)

    def forward(self, x):
        x = self.conv1(self.relu(self.bn1(x)))
        x = self.relu(self.bn2(x))
        if self.dropout is not None:
            x = self.dropout(x)
        out = self.conv_seg(x)
        return out

class AU_SCTHead(nn.Module):
    def __init__(self,in_channels,feats_channels,num_classes,dropout_ratio=0.1, align_corners=False):
        super(AU_SCTHead,self).__init__()
        self.align_corners = align_corners
        self.bn1 = nn.SyncBatchNorm(in_channels)
        self.conv1 = nn.Conv2d(
            in_channels,
            feats_channels,
            kernel_size=3,
            padding=1)
        self.bn2 = nn.SyncBatchNorm(feats_channels)
        self.relu = nn.ReLU()
        if dropout_ratio > 0:
            self.dropout = nn.Dropout2d(dropout_ratio)
        else:
            self.dropout = None
        self.conv_seg = nn.Conv2d(feats_channels, num_classes, kernel_size=1)

    def forward(self, x):
        x = self.conv1(self.relu(self.bn1(x)))
        x = self.relu(self.bn2(x))
        if self.dropout is not None:
            x = self.dropout(x)
        out = self.conv_seg(x)
        return out

class VitGuidanceHead(nn.Module):
    def __init__(self,in_channels,feats_channels,num_classes,dropout_ratio=0.1, align_corners=False,inter_transform_type='linear',base_channels=64):
        super(VitGuidanceHead,self).__init__()
        self.inter_transform_type=inter_transform_type
        self.channelsDefault=[256*4, 320, 512] #decoder,stage3,stage4
        if self.inter_transform_type == 'linear':
            self.feature_transforms =nn.Conv2d(base_channels*4, self.channelsDefault[0], 1)
            self.convstage4= nn.Conv2d(base_channels*8,self.channelsDefault[2],1)
            self.convstage3= nn.Conv2d(base_channels*4,self.channelsDefault[1],1)
        self.cls_seg = None
        self.conv_seg= None
        self.feature_transforms.apply(self._weights_kaiming)
        self.convstage4.apply(self._weights_kaiming)
        self.convstage3.apply(self._weights_kaiming)

        if base_channels==32:
            self.teacher_backbone = MixVisionTransformer(
                structure_type="mit-b2",
                in_channels=3,
                embed_dims=64,
                num_stages=4,
                num_layers=[3, 4, 6, 3],
                num_heads=[1, 2, 5, 8],
                patch_sizes=[7, 3, 3, 3],
                sr_ratios=[8, 4, 2, 1],
                out_indices=(0, 1, 2, 3),
                mlp_ratio=4,
                qkv_bias=True,
                drop_rate=0.0,
                attn_drop_rate=0.0,
                drop_path_rate=0.1,
            )
        else:
            self.teacher_backbone = MixVisionTransformer(
                structure_type="mit-b3",
                in_channels=3,
                embed_dims=64,
                num_stages=4,
                num_layers=[3, 4, 18, 3],
                num_heads=[1, 2, 5, 8],
                patch_sizes=[7, 3, 3, 3],
                sr_ratios=[8, 4, 2, 1],
                out_indices=(0, 1, 2, 3),
                mlp_ratio=4,
                qkv_bias=True,
                drop_rate=0.0,
                attn_drop_rate=0.0,
                drop_path_rate=0.1,
            )

        self.teacher_head = SegformerHead(
            in_channels=[64, 128, 320, 512],
            feats_channels=256,
            dropout_ratio=0.1,
            num_classes=num_classes,
            norm_cfg=dict(type='SyncBN', requires_grad=True),
            align_corners=False,
        )

        self.teacher_backbone.eval()
        self.teacher_head.eval()

    def align_scale(self, stu_feat, tea_feat):
        N, C, H, W = stu_feat.size()
        # normalize student feature
        stu_feat = stu_feat.permute(1, 0, 2, 3).reshape(C, -1)
        stu_mean = stu_feat.mean(dim=-1, keepdim=True)
        stu_std = stu_feat.std(dim=-1, keepdim=True)
        stu_feat = (stu_feat - stu_mean) / (stu_std + 1e-6)
        #
        tea_feat = tea_feat.permute(1, 0, 2, 3).reshape(C, -1)
        tea_mean = tea_feat.mean(dim=-1, keepdim=True)
        tea_std = tea_feat.std(dim=-1, keepdim=True)
        stu_feat = stu_feat * tea_std + tea_mean
        return stu_feat.reshape(C, N, H, W).permute(1, 0, 2, 3)

    def forward(self, inputs):
        x_guidance_feature = inputs
        self.teacher_backbone.eval()
        self.teacher_head.eval()
        with torch.no_grad():
            vitfeature=self.teacher_backbone(x_guidance_feature[0])
            vitfeat,vitseg_logits,feature = self.teacher_head(vitfeature)
            feat_t=[]
            feat_t.append(vitseg_logits)   #shape:4x19x256x128 stride=4
            feat_t.append(vitfeat)   #4x512x256x128 stride=4
            feat_t.append(vitfeature[3])   #4,512,32,16 stride=32
            feat_t.append(vitfeature[2]) #4,320,64,32 stride=16
        stu_channel_align = self.feature_transforms(x_guidance_feature[1][0])
        if vitfeature[0].size(-2)!=stu_channel_align.size(-2) or vitfeature[0].size(-1)!=stu_channel_align.size(-1):
                dsize = (vitfeature[0].size(-2), vitfeature[0].size(-1))
                #feat_t = F.interpolate(feat_t, dsize, mode='bilinear', align_corners=False)
                sampled_decoder_feature = F.interpolate(stu_channel_align, dsize, mode='bilinear', align_corners=False)
        stu_fea = self.align_scale(sampled_decoder_feature, feature)
        stu_decoder_feature = self.teacher_head.fusion_conv(stu_fea)
        stu_decoder_logits = self.teacher_head.conv_seg(stu_decoder_feature)
        if self.inter_transform_type == 'linear':
            feat_s=[]
            feat_s.append(stu_decoder_logits)    #shape:4x19x128x64 stride=8
            feat_s.append(stu_decoder_feature)  #xout,4x256x128x64 stride=8
            feat_s.append(self.convstage4(x_guidance_feature[1][1]))     #x5,  4,512,32,16 stride=32
            feat_s.append(self.convstage3(x_guidance_feature[1][2]))     #x3   4,320,64,32 stride=16
        out = [feat_t,feat_s]
        return out
    def _weights_kaiming(self, m):
        if isinstance(m, nn.Linear):
            nn.init.trunc_normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.constant_(m.bias, val=0)
        elif isinstance(m, (nn.SyncBatchNorm, nn.BatchNorm2d)):
            nn.init.constant_(m.weight, val=1.0)
            nn.init.constant_(m.bias, val=0)
        elif isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight)
            if m.bias is not None:
                nn.init.constant_(m.bias, val=0)



class SegformerHead(nn.Module):
    """The all mlp Head of segformer.

    This head is the implementation of
    `Segformer <https://arxiv.org/abs/2105.15203>` _.

    Args:
        interpolate_mode: The interpolate mode of MLP head upsample operation.
            Default: 'bilinear'.
    """

    def __init__(self,in_channels,feats_channels,num_classes,norm_cfg,dropout_ratio=0.1, align_corners=False,):
        super(SegformerHead,self).__init__()
        self.align_corners = align_corners
        num_inputs = len(in_channels)

        self.convs = nn.ModuleList()
        for i in range(num_inputs):
            self.convs.append(
                nn.Sequential(
                nn.Conv2d(in_channels[i], feats_channels, kernel_size=1, stride=1, padding=0, bias=False),
                nn.SyncBatchNorm(feats_channels),
                nn.ReLU()
            ))

        self.fusion_conv = nn.Sequential(
            nn.Conv2d(feats_channels * len(self.convs),feats_channels, kernel_size=1, stride=1, padding=0, bias=False),
            nn.SyncBatchNorm(feats_channels),
            nn.ReLU()
        )
        if dropout_ratio > 0:
            self.dropout = nn.Dropout2d(dropout_ratio)
        else:
            self.dropout = None
        self.conv_seg = nn.Conv2d(feats_channels, num_classes, kernel_size=1)
    def forward(self, inputs):
        # Receive 4 stage backbone feature map: 1/4, 1/8, 1/16, 1/32
        outs = []
        for idx in range(len(inputs)):
            x = inputs[idx]
            conv = self.convs[idx]
            outs.append(
                F.interpolate(
                    conv(x),
                    inputs[0].shape[2:],
                    mode='bilinear',
                    align_corners=self.align_corners))
        feature = torch.cat(outs, dim=1)
        out = self.fusion_conv(feature)
        if self.dropout is not None:
            out = self.dropout(out)
        logits = self.conv_seg(out)

        return out,logits,feature
